import numpy as np

class UAV:
    def __init__(self, grid_size, id, charging_stations, num_position_action, num_energy_action):
        # 无人机基础信息
        self.id = id
        self.speed = 25
        
        # 无人机能量信息
        self.energy = 100
        self.has_charge = False
        self.charging_stations = charging_stations
        self.num_energy_action = num_energy_action
        self.epsilon = 0.1
        
        # 无人机的位置信息
        self.grid_size = grid_size
        self.position = (self.id*50, self.id*30)
        self.num_position_action = num_position_action
        
        # 无人机历史信息
        self.history_position = []
        self.history_length = 2
        
    def execute_cruise_action(self, action):
        """执行巡航动作 - 优化版本"""
        # 转换为整数列表
        if isinstance(action, list):
            pass
        else:
            action = action.astype(int).flatten().tolist()
        
        # 当前坐标
        x, y = self.position
        distance_per_time = self.speed
        reward = 0

        # 根据动作更新位置
        if action == [1,0,0,0]:  # 向上移动
            y += distance_per_time
        elif action == [0,1,0,0]:  # 向下移动
            y -= distance_per_time
        elif action == [0,0,1,0]:  # 向左移动
            x -= distance_per_time
        elif action == [0,0,0,1]:  # 向右移动
            x += distance_per_time
        # 停留位置不变
        
        # 检查是否超出边界并限制 - 减轻边界惩罚
        if x < 0:
            x = 0
            reward -= 20  # 减少边界惩罚从-100到-20
        elif x >= self.grid_size:
            x = self.grid_size - 1
            reward -= 20
        
        if y < 0:
            y = 0
            reward -= 20
        elif y >= self.grid_size:
            y = self.grid_size - 1
            reward -= 20
        
        # 更新位置
        self.position = (x, y)

        # 避免重复访问同一位置 - 减轻惩罚
        if self.position in self.history_position:
            reward -= 2  # 减少重复访问惩罚从-10到-2
        
        # 更新历史路径
        self.history_position.append(self.position)
        if len(self.history_position) > self.history_length:
            self.history_position.pop(0)
        
        # 执行电量消耗
        self.reduce_energy(1)
        return reward, self.history_position

    def reduce_energy(self, amount):
        """扣除电量"""
        self.energy -= amount
        if self.energy < 0:
            self.energy = 0
            
    def get_state(self):
        """获取状态，包括位置信息"""
        return np.array([self.position[0], self.position[1]])
        
    def get_extended_state(self, target_positions, other_uav_positions):
        """获取扩展状态，包括多个目标和其他无人机的相对位置"""
        state = [self.position[0], self.position[1]]  # 自身位置
        
        # 添加到所有目标的相对位置（最多2个目标）
        for i, target_pos in enumerate(target_positions[:2]):
            rel_target_x = target_pos[0] - self.position[0]
            rel_target_y = target_pos[1] - self.position[1]
            state.extend([rel_target_x, rel_target_y])
        
        # 如果目标数量不足2个，用0填充
        while len(state) < 6:  # 2(自身) + 4(两个目标)
            state.append(0)
        
        # 添加到其他无人机的相对位置（最多考虑最近的2个）
        distances = []
        for other_pos in other_uav_positions:
            if other_pos != self.position:
                dist = np.linalg.norm(np.array(other_pos) - np.array(self.position))
                distances.append((dist, other_pos))
        
        distances.sort(key=lambda x: x[0])
        
        # 添加最近的2个无人机的相对位置
        for i in range(min(2, len(distances))):
            _, other_pos = distances[i]
            rel_x = other_pos[0] - self.position[0]
            rel_y = other_pos[1] - self.position[1]
            state.extend([rel_x, rel_y])
        
        # 如果无人机数量不足，用0填充
        while len(state) < 10:  # 2(自身) + 4(目标) + 4(其他无人机)
            state.append(0)
            
        # 添加能量状态
        state.append(self.energy / 100.0)  # 归一化能量
        
        # 添加到最近目标的距离（归一化）
        min_target_dist = float('inf')
        for target_pos in target_positions:
            dist = np.linalg.norm(np.array(target_pos) - np.array(self.position))
            min_target_dist = min(min_target_dist, dist)
        
        state.append(min_target_dist / 200.0)  # 归一化距离
        
        return np.array(state[:12])  # 限制状态维度为12


